# step2_filter_standardize.py
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler

print("=== 步骤2: 数据筛选和标准化 ===")

# 定义列名
column_names = ['class', 'alcohol', 'malic_acid', 'ash', 'alcalinity_of_ash', 
                'magnesium', 'total_phenols', 'flavanoids', 'nonflavanoid_phenols',
                'proanthocyanins', 'color_intensity', 'hue', 'od280/od315', 'proline']

try:
    # 读取数据
    print("正在读取数据...")
    wine_data = pd.read_csv('wine.data', header=None, names=column_names)
    print(f"原始数据读取成功！形状: {wine_data.shape}")
    
    # 1. 筛选类别1和2的数据
    print("\n筛选类别1和2的数据...")
    filtered_data = wine_data[wine_data['class'].isin([1, 2])].copy()
    
    print(f"筛选后的数据形状: {filtered_data.shape}")
    print("筛选后各类别样本数量:")
    class_counts = filtered_data['class'].value_counts().sort_index()
    for class_label, count in class_counts.items():
        print(f"  类别 {class_label}: {count} 个样本")
    
    # 2. 准备特征和标签
    print("\n准备特征和标签...")
    X = filtered_data.drop('class', axis=1)  # 特征（去掉类别列）
    y = filtered_data['class']               # 标签（类别列）
    
    print(f"特征矩阵 X 的形状: {X.shape}")
    print(f"标签向量 y 的形状: {y.shape}")
    print(f"特征名称: {list(X.columns)}")
    
    # 3. 数据标准化
    print("\n数据标准化处理...")
    print("标准化前 - 前3个样本的第一个特征值:")
    print(f"  样本1: {X.iloc[0, 0]:.4f}")
    print(f"  样本2: {X.iloc[1, 0]:.4f}")
    print(f"  样本3: {X.iloc[2, 0]:.4f}")
    
    # 使用StandardScaler进行标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    print("标准化后 - 前3个样本的第一个特征值:")
    print(f"  样本1: {X_scaled[0, 0]:.4f}")
    print(f"  样本2: {X_scaled[1, 0]:.4f}")
    print(f"  样本3: {X_scaled[2, 0]:.4f}")
    
    # 检查标准化效果
    print(f"\n标准化完成！")
    print(f"标准化后的特征矩阵形状: {X_scaled.shape}")
    print(f"标准化后数据的均值: {np.mean(X_scaled, axis=0)[:3]:.4f}...")  # 显示前3个特征的均值
    print(f"标准化后数据的标准差: {np.std(X_scaled, axis=0)[:3]:.4f}...") # 显示前3个特征的标准差
    
    # 保存处理后的数据
    processed_data = pd.DataFrame(X_scaled, columns=X.columns)
    processed_data['class'] = y.values
    processed_data.to_csv('wine_processed.csv', index=False)
    print(f"\n处理后的数据已保存为: wine_processed.csv")
    
    # 显示一些统计信息
    print("\n数据统计信息:")
    print(f"总样本数: {len(X_scaled)}")
    print(f"特征数: {X_scaled.shape[1]}")
    print(f"类别1样本数: {sum(y == 1)}")
    print(f"类别2样本数: {sum(y == 2)}")
    
except Exception as e:
    print(f"出错: {e}")

input("\n按 Enter 键继续下一步...")